package com.game.sharding.config;

import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import lombok.Generated;
import org.apache.commons.lang3.StringUtils;
import org.apache.shardingsphere.infra.context.transaction.TransactionConnectionContext;
import org.apache.shardingsphere.readwritesplitting.spi.ReadQueryLoadBalanceAlgorithm;

import java.util.List;
import java.util.Properties;

public class CustomLoadBalanceAlgorithm implements ReadQueryLoadBalanceAlgorithm {
    private Properties props;

    public CustomLoadBalanceAlgorithm() {

    }

    @Override
    public void init(Properties props) {
        this.props = props;
    }

    /**
     * 获取数据源
     *
     * @param name  数据源名称（ShardingJDBC使用的）
     * @param writeDataSourceName 写数据源名称
     * @param readDataSourceNames 所有配置的复杂均衡中读数据源名称
     * @param context 事务上下文对象，可以获取context.isInTransaction() 判断是否需要事务，可通过这个来判断是否使用 写数据源
     * @return java.lang.String
     */
    @Override
    public String getDataSource(String name, String writeDataSourceName, List<String> readDataSourceNames, TransactionConnectionContext context) {
        // 获取当前MybatisPlus指定的数据源
        String dsKey = DynamicDataSourceContextHolder.peek();
        if (StringUtils.isNotBlank(dsKey)) {
            if (writeDataSourceName.equals(dsKey)) {
                return dsKey;
            }
            if (readDataSourceNames.contains(dsKey)) {
                return dsKey;
            }
            throw new RuntimeException("@DS 配置错误，当前数据源[" + dsKey + "]不在SharingJDBC数据源列表[" + readDataSourceNames + "]中");
        }
        return writeDataSourceName;
    }

    @Override
    public String getType() {
        return "CUSTOM";
    }

    @Override
    public boolean isDefault() {
        return true;
    }

    @Override
    @Generated
    public Properties getProps() {
        return this.props;
    }
}